# --- 导入必要的库 ---
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
from manim import *

# --- Case 2: gamma_c > 0 (稳定螺线点) ---


# plot(a_star_damped, b_star_damped, 'o', markersize=8, markerfacecolor='red', markeredgecolor='black', markeredgewidth=1.5)

# 绘制示例轨迹
# initial_conditions = [[a_star_damped*1.5, b_star_damped*1.5], [a_star_damped*0.5, b_star_damped*0.5]]
# t_span = np.linspace(0, 50, 500)
# colors = ['red', 'orange']
# for i, y0 in enumerate(initial_conditions):
#     sol = odeint(lotka_volterra_ode, y0, t_span, args=(lambda_rate, mu_rate, gamma_p, gamma_c))
#     ax.plot(sol[:, 0], sol[:, 1], color=colors[i], lw=2, alpha=0.8)

# ax.set_xlabel('a', fontsize=14)
# ax.set_ylabel('b', fontsize=14)
# ax.set_xlim(0, a_max)
# ax.set_ylim(0, b_max)

# fig.suptitle('Phase Portraits', fontsize=20, color='darkgreen', y=1.0)
# plt.tight_layout()
# plt.show()

class DrawPhasePortraits(Scene):
    def construct(self):
        a_max, b_max = 2.0, 2.0

        def lotka_volterra_ode(y, t, lambda_rate, mu_rate, gamma_p, gamma_c):
            a, b = y
            da_dt = mu_rate * a - gamma_p * a * b
            db_dt = -lambda_rate * b + gamma_p * a * b - gamma_c * b * b
            return [da_dt, db_dt]

        # --- 设置参数 ---
        lambda_rate = 1.0
        mu_rate = 1.0
        gamma_p = 1.0
        gamma_c = 0.001

        # 绘制不动点
        b_star_damped = lambda_rate / gamma_p
        a_star_damped = (mu_rate - gamma_c * b_star_damped) / gamma_p
        axes = Axes(x_range=[0, a_max, a_max/4], y_range=[0, b_max, b_max/4])
        axes_labels = axes.get_axis_labels(x_label="Fox", y_label="Rabbit")
        fix_point = Dot(axes.c2p(a_star_damped, b_star_damped),color=RED)
        self.play(Create(axes), Write(axes_labels))
        self.play(Create(fix_point))
        self.wait(1)

        initial_condition = [a_star_damped*0.5, b_star_damped*0.5]
        t_span = np.linspace(0, 3000, 6000)
        sol = odeint(lotka_volterra_ode, initial_condition, t_span, args=(lambda_rate, mu_rate, gamma_p, gamma_c))
        points_positions = [axes.c2p(x,y) for x, y in sol]
        dots = VGroup(*[
            Dot(point_pos, color=RED, radius=0.03, fill_opacity=0.8)
            for point_pos in points_positions
        ])
        
        # draw the inital 10 dots
        dots = dots[0:200]
        self.play(
            Succession(*[Create(dot) for dot in dots], run_time=1)
        )
        def update_dots(dots, alpha):
            for i, dot in enumerate(dots):
                dot.move_to(points_positions[int(1900*alpha)+i])

        self.play(
            UpdateFromAlphaFunc(dots, update_dots),
            run_time = 15,
            rate_func = linear,
        )
        
        self.wait(1)